package com.bjbn.aesbcommonapi.consumer;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.bjbn.aesbcommonapi.config.RocketMQConfig;
import com.bjbn.aesbcommonapi.service.MessageTransferService;
import lombok.RequiredArgsConstructor;
import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyContext;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyStatus;
import org.apache.rocketmq.client.consumer.listener.MessageListenerConcurrently;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.common.consumer.ConsumeFromWhere;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.acl.common.AclClientRPCHook;
import org.apache.rocketmq.remoting.RPCHook;
import org.apache.rocketmq.common.protocol.heartbeat.MessageModel;
import org.apache.rocketmq.acl.common.SessionCredentials;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.CommandLineRunner;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;
import javax.annotation.Resource;

import com.bjbn.aesbcommonapi.utils.MsgFormatUtil;

@Component
@RequiredArgsConstructor
public class RocketMQConsumer implements CommandLineRunner {
    private static final Logger logger = LoggerFactory.getLogger(RocketMQConsumer.class);
    private final RocketMQConfig rocketMQConfig;
    private final MessageTransferService transferService;

    // 注入Spring配置的异步执行器
    @Resource(name = "asyncTaskExecutor")
    private Executor asyncExecutor;

    // JSON解析缓存，避免重复解析
    private final ConcurrentHashMap<String, JSONObject> jsonCache = new ConcurrentHashMap<>();

    @Override
    public void run(String... args) throws Exception {
        RPCHook rpcHook = new AclClientRPCHook(new SessionCredentials(
                rocketMQConfig.getAccessKey(), rocketMQConfig.getSecretKey()));
        if (rocketMQConfig.getGroups() == null) {
            logger.warn("未配置任何消费者组");
            return;
        }
        for (RocketMQConfig.GroupConfig group : rocketMQConfig.getGroups()) {
            DefaultMQPushConsumer consumer = new DefaultMQPushConsumer(
                    null,
                    group.getConsumerGroup(),
                    rpcHook,
                    new org.apache.rocketmq.client.consumer.rebalance.AllocateMessageQueueAveragely(),
                    true,
                    null
            );
            consumer.setNamesrvAddr(rocketMQConfig.getNamesrvAddr());
            consumer.setConsumeFromWhere(ConsumeFromWhere.CONSUME_FROM_LAST_OFFSET);
            consumer.setMessageModel(MessageModel.CLUSTERING);

            // 设置消费者参数，提升性能
            consumer.setConsumeThreadMin(10);
            consumer.setConsumeThreadMax(20);
            consumer.setPullBatchSize(32);
            consumer.setConsumeMessageBatchMaxSize(10);

            if (Boolean.TRUE.equals(rocketMQConfig.getUseTls())) {
                consumer.setUseTLS(true);
            }
            // 订阅该 group 下所有 topic
            for (RocketMQConfig.TopicMapping mapping : group.getTopics()) {
                consumer.subscribe(mapping.getSource(), "*");
            }
            consumer.registerMessageListener(new MessageListenerConcurrently() {
                @Override
                public ConsumeConcurrentlyStatus consumeMessage(List<MessageExt> msgs, ConsumeConcurrentlyContext context) {
                    // 异步处理消息，避免阻塞消费者线程
                    CompletableFuture.runAsync(() -> {
                        processMessagesAsync(msgs, group);
                    }, asyncExecutor);

                    return ConsumeConcurrentlyStatus.CONSUME_SUCCESS;
                }
            });
            try {
                consumer.start();
                logger.info("RocketMQ消费者已启动，group: {}, topics: {}", group.getConsumerGroup(), group.getTopics().stream().map(RocketMQConfig.TopicMapping::getSource).toArray());
            } catch (MQClientException e) {
                logger.error("启动RocketMQ消费者失败，group: {}", group.getConsumerGroup(), e);
            }
        }
    }

    /**
     * 异步处理消息列表
     */
    @Async
    public void processMessagesAsync(List<MessageExt> msgs, RocketMQConfig.GroupConfig group) {
        for (MessageExt msg : msgs) {
            try {
                // 查找该 group 下的 topic 映射
                String kafkaTopic = null;
                String resourceTopic = null;
                for (RocketMQConfig.TopicMapping mapping : group.getTopics()) {
                    if (mapping.getSource().equals(msg.getTopic())) {
                        kafkaTopic = mapping.getTarget();
                        resourceTopic = msg.getTopic();
                        break;
                    }
                }
                if (kafkaTopic != null) {
                    final String finalKafkaTopic = kafkaTopic;
                    final String finalResourceTopic = resourceTopic;
                    logger.debug("接收到 RocketMQ消息， topic: {}，消息ID: {}", kafkaTopic, msg.getMsgId());
                    // 异步处理消息转发到kafka，通过kafka存储到原始数据到es
                    CompletableFuture.runAsync(() -> {
                        transferService.transfer(finalResourceTopic, msg);
                    }, asyncExecutor);

                    // 异步处理消息发送到目标主题
                    CompletableFuture.runAsync(() -> {
                        sendMsg2OwnTopic(finalKafkaTopic, new String(msg.getBody()));
                    }, asyncExecutor);
                } else {
                    logger.warn("未找到对应Kafka topic, RocketMQ topic: {}", msg.getTopic());
                }
            } catch (Exception e) {
                logger.error("处理消息失败，topic: {}, messageId: {}", msg.getTopic(), msg.getMsgId(), e);
            }
        }
    }

    /**
     * 优化后的消息发送方法
     */
    private void sendMsg2OwnTopic(String targetTopic, String payload) {
        try {
            JSONObject jsonObject = JSONObject.parseObject(payload);

            if (MsgFormatUtil.isStandardFormat(jsonObject)) {
                // 标准格式消息，需要判断data是数组还是对象
                String key = "data";
                Object value = jsonObject.get(key);
                if (value instanceof JSONArray) {
                    JSONArray array = jsonObject.getJSONArray(key);

                    for (int i = 0; i < array.size(); i++) {
                        JSONObject jsonObj = array.getJSONObject(i);
                        transferService.transferBody(targetTopic, jsonObj);
                    }

                } else if (value instanceof JSONObject) {
                    JSONObject obj = jsonObject.getJSONObject(key);
                    transferService.transferBody(targetTopic, obj);
                }
            } else {
                // 其他格式的消息 直接发送
                transferService.transferBody(targetTopic, jsonObject);
            }

        } catch (Exception e) {
            logger.error("发送消息到目标主题失败，topic: {}, payload: {}", targetTopic, payload, e);
        }
    }


}